import pandas as pd  # 导入pandas库对数据进行读取
from sklearn.feature_extraction.text import TfidfVectorizer
from pathlib import Path
from sklearn.metrics import roc_auc_score
from sklearn.preprocessing import LabelEncoder, LabelBinarizer, StandardScaler  # 导入标签编码器
from sklearn.model_selection import train_test_split, GridSearchCV  # 导入数据划分器
from sklearn.linear_model import LogisticRegression  # 导入逻辑回归模型做分类
from xgboost import XGBClassifier
import joblib  # 导入模型保存模块
import pickle



train_file = "D:/code/datasets/train_v2_drcat_02/train_v2_drcat_02.csv"
# 加载数据
df = pd.read_csv(train_file, usecols=['text', 'label'])
# val_df = pd.read_csv('val.csv', usecols=['text', 'label'])
"""
print(df.info())
 0   text    44868 non-null  object
 1   label   44868 non-null  int64
"""

# 对标签进行编码
label_encoder = LabelEncoder()
df['label'] = label_encoder.fit_transform(df['label'])
print(df)
# 文本向量化
vectorizer = TfidfVectorizer()
X = vectorizer.fit_transform(df['text'])

# 保存 vectorizer
joblib.dump(vectorizer, str(Path("D:/code/models") / "vectorizer.pkl"))

y = df['label']

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
# # 训练模型
# logic_clf = LogisticRegression(max_iter=1000)
# logic_clf.fit(X_train, y_train)
# XGBClassifier 方式训练
xgb_clf = XGBClassifier()
xgb_clf.fit(X_train, y_train)

# 预测验证集 方法1
y_pred = xgb_clf.predict_proba(X_test)[:, 1]
print(y_pred)
# 计算AUC
auc = roc_auc_score(y_test, y_pred)
print(f'logic AUC {auc}')

joblib.dump(xgb_clf, str(Path("D:/code/models") / "logist_regression_model1.pkl"))

val_file = "D:/code/datasets/train_v2_drcat_02/val.csv"
# 加载数据
df_val = pd.read_csv(val_file, usecols=['text'])

# 加载 vectorizer
vectorizer = joblib.load(str(Path("D:/code/models") / "vectorizer.pkl"))
# 使用训练集上的 TfidfVectorizer 来转换验证集数据
X_val = vectorizer.transform(df_val['text'])  # 注意这里使用 transform 而不是 fit_transform

# 加载模型
model = joblib.load(str(Path("D:/code/models") / "logist_regression_model1.pkl"))
# # 预测验证集
# y_pred_val = model1.predict_proba(X_val)[:, 1]
# print(y_pred_val)
y_pred_val = model.predict(X_val)
pre = pd.DataFrame(y_pred_val,columns=['label'])
# 指定数据集合的位置
text =  df_val["text"]
submission = pd.concat([text,pre],axis=1)
submission.to_csv("val.csv",index=False)